package com.koicarp.agent.provider;

import java.io.IOException;
import java.util.concurrent.ExecutorService;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import com.shsajt.ai.common.ApiResult;
import com.shsajt.ai.provider.constant.ProviderErrMsgConstant;

/**
 * @ClassName:  StreamEmitter   
 * @Description: 流式emitter  
 * @author: liutao
 * @date:   2025年3月6日 下午5:04:43    
 */
public class StreamEmitter {

    private final SseEmitter emitter;

    private final Logger log = LoggerFactory.getLogger(StreamEmitter.class);
    
    public StreamEmitter() {
        emitter = new SseEmitter(5 * 60 * 1000L);
    }

    public SseEmitter get() {
        return emitter;
    }

    public SseEmitter streaming(final ExecutorService executor, Runnable func) {
        emitter.onCompletion(executor::shutdownNow);
        emitter.onError(e -> {
        	log.error(ProviderErrMsgConstant.SSE_ERROR, e.getMessage());
            executor.shutdownNow();
        });
        emitter.onTimeout(() -> {
        	log.error(ProviderErrMsgConstant.SSE_TIMEOUT);
            emitter.complete();
            executor.shutdownNow();
        });
        
        executor.execute(() -> {
            try {
                func.run();
            } catch (Exception e) {
                emitter.completeWithError(e);
                Thread.currentThread().interrupt();
            } finally {
                if (!executor.isShutdown()) {
                    executor.shutdownNow();
                }
            }
        });
        return emitter;
    }

    public void send(Object obj) {
        try {
            emitter.send(obj);
        } catch (IOException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    public void complete() {
        emitter.complete();
    }

    public void error(ApiResult result) {
        try {
            emitter.send(result);
            emitter.complete();
        } catch (IOException e) {
            throw new RuntimeException(e.getMessage());
        }
    }

}
